#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TCP 端口转发工具
用法: python tcp_forward.py <listen_port> <target_host> <target_port>
示例: python tcp_forward.py 19893 127.0.0.1 9222
"""

import socket
import threading
import sys
import time

BUFSIZE = 32 * 1024  # 单次最大缓冲 32 KiB


def log(msg):
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}")


def relay(sock_in: socket.socket, sock_out: socket.socket):
    """单向转发数据，直到 EOF"""
    try:
        while True:
            data = sock_in.recv(BUFSIZE)
            if not data:          # 对端关闭
                break
            sock_out.sendall(data)
    except OSError:
        pass
    finally:
        sock_in.close()
        sock_out.close()


def handle_client(client_conn: socket.socket, target_host: str, target_port: int):
    """接受一个客户端连接，并把它桥接到目标端口"""
    try:
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        remote.connect((target_host, target_port))
    except Exception as e:
        log(f"无法连接到目标 {target_host}:{target_port} -> {e}")
        client_conn.close()
        return

    # 启动两条线程，双向转发
    threading.Thread(target=relay, args=(client_conn, remote), daemon=True).start()
    threading.Thread(target=relay, args=(remote, client_conn), daemon=True).start()


def main():
    if len(sys.argv) not in (3, 4):
        print("用法: python tcp_forward.py <listen_port> <target_host> <target_port>")
        print("本机转发示例: python tcp_forward.py 19893 127.0.0.1 9222")
        sys.exit(1)

    listen_port = int(sys.argv[1])
    target_host = sys.argv[2] if len(sys.argv) == 4 else "127.0.0.1"
    target_port = int(sys.argv[-1])

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server.bind(("", listen_port))
        server.listen(100)
        log(f"监听 *:{listen_port}  ->  {target_host}:{target_port}")
    except OSError as e:
        log(f"无法监听端口 {listen_port}: {e}")
        sys.exit(2)

    try:
        while True:
            conn, addr = server.accept()
            log(f"新连接 {addr}")
            threading.Thread(
                target=handle_client,
                args=(conn, target_host, target_port),
                daemon=True
            ).start()
    except KeyboardInterrupt:
        log("收到 Ctrl-C，正在退出…")
    finally:
        server.close()


if __name__ == "__main__":
    main()
